import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

def load_data():
    """加载2024实际数据和2028预测数据"""
    try:
        # 读取2024年数据
        df_2024 = pd.read_csv('summerOly_medal_counts.csv')
        # 打印列名以检查
        print("2024数据列名:", df_2024.columns.tolist())
        df_2024 = df_2024[df_2024['Year'] == 2024].copy()
        
        # 读取2028年预测数据
        df_2028 = pd.read_csv('2028_olympics_predictions.csv')
        # 打印列名以检查
        print("2028预测数据列名:", df_2028.columns.tolist())
        df_2028.set_index('Unnamed: 0', inplace=True)
        
        return df_2024, df_2028
    except Exception as e:
        print(f"数据加载错误: {str(e)}")
        return None, None

def calculate_progress(row):
    """计算进步程度"""
    try:
        # 计算总奖牌变化百分比
        total_change_pct = ((row['predicted_total'] - row['Total_2024']) / row['Total_2024'] * 100 
                           if row['Total_2024'] > 0 else 0)
        
        # 计算金牌变化百分比
        gold_change_pct = ((row['predicted_gold'] - row['Gold_2024']) / row['Gold_2024'] * 100 
                          if row['Gold_2024'] > 0 else 0)
        
        # 返回英文结果
        if total_change_pct > 10 or gold_change_pct > 10:
            return 'Significant Progress'
        elif total_change_pct > 5 or gold_change_pct > 5:
            return 'Slight Progress'
        elif total_change_pct < -10 or gold_change_pct < -10:
            return 'Significant Decline'
        elif total_change_pct < -5 or gold_change_pct < -5:
            return 'Slight Decline'
        else:
            return 'Stable'
    except Exception as e:
        print(f"Progress calculation error: {str(e)}")
        return 'Calculation Error'

def visualize_changes(comparison):
    """可视化奖牌变化"""
    # 设置字体
    plt.rcParams['font.sans-serif'] = ['Arial']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 创建保存目录
    Path('predictions/plots').mkdir(parents=True, exist_ok=True)
    
    # 1. 进步程度分布饼图
    plt.figure(figsize=(10, 8))
    
    # 直接使用英文标签的数据
    progress_counts = comparison['progress'].value_counts()  # 因为calculate_progress已经返回英文结果
    
    # 设置颜色
    colors = ['#2ecc71', '#3498db', '#95a5a6', '#e74c3c', '#e67e22']
    
    # 创建饼图
    plt.pie(progress_counts, 
            labels=progress_counts.index,
            colors=colors,
            autopct='%1.1f%%',
            textprops={'fontsize': 10})
    
    plt.title('Distribution of Progress Levels', fontsize=14, pad=20)
    
    # 添加图例
    plt.legend(progress_counts.index,
              title='Progress Levels',
              loc='center left',
              bbox_to_anchor=(1, 0, 0.5, 1))
    
    # 保存图片
    plt.savefig('predictions/plots/progress_distribution.png', 
                dpi=300, 
                bbox_inches='tight',
                pad_inches=0.2)
    plt.close()
    
    # 2. 变化幅度热力图
    plt.figure(figsize=(12, 6))
    changes_data = comparison[['total_change_pct', 'gold_change_pct']].head(20)
    changes_data.columns = ['Total Medals Change %', 'Gold Medals Change %']
    
    sns.heatmap(changes_data.T, annot=True, cmap='RdYlBu', center=0)
    plt.title('TOP 20 Countries Medal Change Percentage', fontsize=12)
    plt.savefig('predictions/plots/change_heatmap.png', dpi=300, bbox_inches='tight')
    plt.close()

def generate_detailed_report(comparison):
    """生成详细的分析报告"""
    # 创建保存目录
    Path('predictions').mkdir(exist_ok=True)
    
    # 生成HTML格式的详细报告
    html_content = """
    <html>
    <head>
        <title>2024-2028奥运会奖牌变化分析报告</title>
        <style>
            body { font-family: Arial, sans-serif; margin: 20px; }
            table { border-collapse: collapse; width: 100%; }
            th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
            th { background-color: #f2f2f2; }
            .progress { color: green; }
            .regress { color: red; }
        </style>
    </head>
    <body>
        <h1>2024-2028奥运会奖牌变化分析报告</h1>
    """
    
    # 添加总体统计
    progress_stats = comparison['progress'].value_counts()
    html_content += "<h2>总体变化情况</h2>"
    html_content += "<ul>"
    for category, count in progress_stats.items():
        html_content += f"<li>{category}: {count}个国家</li>"
    html_content += "</ul>"
    
    # 添加详细数据表格
    html_content += "<h2>各国详细变化情况</h2>"
    html_content += "<table>"
    html_content += """
        <tr>
            <th>国家</th>
            <th>2024总奖牌</th>
            <th>2028预测</th>
            <th>变化数量</th>
            <th>变化百分比</th>
            <th>进步程度</th>
        </tr>
    """
    
    # 按变化幅度排序
    sorted_data = comparison.sort_values('total_change', ascending=False)
    for idx, row in sorted_data.iterrows():
        color_class = 'progress' if row['total_change'] > 0 else 'regress'
        html_content += f"""
            <tr>
                <td>{idx}</td>
                <td>{row['Total_2024']}</td>
                <td>{row['predicted_total']}</td>
                <td class="{color_class}">{row['total_change']:+.1f}</td>
                <td class="{color_class}">{row['total_change_pct']:+.1f}%</td>
                <td>{row['progress']}</td>
            </tr>
        """
    
    html_content += "</table></body></html>"
    
    # 保存HTML报告
    with open('predictions/medal_comparison_report.html', 'w', encoding='utf-8') as f:
        f.write(html_content)

def compare_results():
    """比较2024和2028的结果"""
    df_2024, df_2028 = load_data()
    if df_2024 is None or df_2028 is None:
        return None
    
    # 创建比较结果DataFrame
    comparison = pd.DataFrame()
    
    try:
        # 假设国家列名可能是'country'或'NOC'
        country_column = 'country' if 'country' in df_2024.columns else 'NOC'
        
        # 合并2024和2028数据
        for index, row in df_2028.iterrows():
            country_2024 = df_2024[df_2024[country_column] == index]
            if len(country_2024) > 0:
                comparison.loc[index, 'Total_2024'] = country_2024['Total'].iloc[0]
                comparison.loc[index, 'Gold_2024'] = country_2024['Gold'].iloc[0]
                comparison.loc[index, 'predicted_total'] = row['predicted_total']
                comparison.loc[index, 'predicted_gold'] = row['predicted_gold']
                
        # 计算变化
        comparison['total_change'] = comparison['predicted_total'] - comparison['Total_2024']
        comparison['gold_change'] = comparison['predicted_gold'] - comparison['Gold_2024']
        comparison['total_change_pct'] = (comparison['total_change'] / comparison['Total_2024'] * 100).round(1)
        comparison['gold_change_pct'] = (comparison['gold_change'] / comparison['Gold_2024'] * 100).round(1)
        
        # 评估进步情况
        comparison['progress'] = comparison.apply(calculate_progress, axis=1)
        
        # 生成可视化
        visualize_changes(comparison)
        
        # 生成详细报告
        generate_detailed_report(comparison)
        
        # 保存结果
        comparison.to_csv('predictions/medal_comparison_2024_2028.csv')
        
        return comparison
    
    except Exception as e:
        print(f"比较过程错误: {str(e)}")
        return None

def print_analysis(comparison):
    """打印分析结果"""
    if comparison is None:
        print("无法进行分析，数据处理出错")
        return
        
    print("\n=== 2024-2028奥运会奖牌变化分析 ===\n")
    
    try:
        # 进步最显著的国家
        print("进步最显著的国家（总奖牌增加）:")
        progress = comparison[comparison['total_change'] > 0].head(10)
        for idx, row in progress.iterrows():
            print(f"{idx}: +{row['total_change']}枚 (总奖牌增长{row['total_change_pct']}%)")
        
        print("\n退步最显著的国家（总奖牌减少）:")
        regress = comparison[comparison['total_change'] < 0].head(10)
        for idx, row in regress.iterrows():
            print(f"{idx}: {row['total_change']}枚 (总奖牌减少{abs(row['total_change_pct'])}%)")
        
        # 统计各类变化的国家数量
        progress_stats = comparison['progress'].value_counts()
        print("\n变化情况统计:")
        for category, count in progress_stats.items():
            print(f"{category}: {count}个国家")
            
        print("\n详细的分析报告已生成到 predictions/medal_comparison_report.html")
        print("可视化结果已保存到 predictions/plots")
            
    except Exception as e:
        print(f"分析打印错误: {str(e)}")

if __name__ == "__main__":
    comparison = compare_results()
    print_analysis(comparison) 